import gradio as gr
import pandas as pd
import os
from excel_analyzer import ExcelAnalyzer
import json
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime

class ChatExcelUI:
    def __init__(self):
        self.secret_id = "AKIDRmvrbAWfbP2UmrjNo6CHOSVREsbHqHWo"
        self.secret_key = "IVZaQ3iGsMySjmcOwLYbUBeuxmOmwtd9"
        self.analyzer = ExcelAnalyzer(self.secret_id, self.secret_key)
        self.current_df = None
        self.chat_history = []
        
        # 设置matplotlib中文支持
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
        plt.rcParams['axes.unicode_minus'] = False     # 用来正常显示负号
        
        # 设置seaborn样式
        sns.set_theme(font='SimHei', font_scale=1.2)
        
        # 使用 analyzer 的图表类型
        self.available_charts = self.analyzer.available_charts
        
    def load_excel(self, file):
        """加载Excel文件并转换为DuckDB表"""
        try:
            if file is None:
                return "请选择文件", gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[])
            
            file_path = file.name
            file_ext = os.path.splitext(file_path)[1].lower()
            
            # 处理图片文件
            if file_ext in ['.png', '.jpg', '.jpeg', '.bmp', '.gif']:
                excel_path = os.path.splitext(file_path)[0] + "_ocr.xlsx"
                success = self.analyzer.image_to_excel(file_path, excel_path)
                if not success:
                    return "图片OCR转换失败", gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[])
                file_path = excel_path
            
            # 读取Excel
            self.current_df = pd.read_excel(file_path)
            # 同时更新 analyzer 的数据
            self.analyzer.current_df = self.current_df
            
            # 生成详细的数据预览
            preview = "## 数据概览\n\n"
            
            # 基本信息
            preview += "### 基本信息\n"
            preview += f"- 总行数：{len(self.current_df)}\n"
            preview += f"- 总列数：{len(self.current_df.columns)}\n"
            preview += f"- 列名：{', '.join(self.current_df.columns)}\n\n"
            
            # 数据类型信息
            preview += "### 数据类型\n"
            preview += "```\n"
            preview += self.current_df.dtypes.to_string()
            preview += "\n```\n\n"
            
            # 完整数据展示
            preview += "### 完整数据\n"
            # 设置pandas显示选项
            pd.set_option('display.max_rows', None)
            pd.set_option('display.max_columns', None)
            pd.set_option('display.width', None)
            preview += self.current_df.to_markdown(index=False)
            
            # 重置pandas显示选项
            pd.reset_option('display.max_rows')
            pd.reset_option('display.max_columns')
            pd.reset_option('display.width')
            
            # 获取列名列表
            columns = list(self.current_df.columns)
            
            # 根据数据类型分类
            numeric_cols = self.current_df.select_dtypes(include=['number']).columns.tolist()
            categorical_cols = self.current_df.select_dtypes(include=['object']).columns.tolist()
            
            # 生成列名提示
            preview += "\n### 数据列类型\n"
            if numeric_cols:
                preview += "**数值型列：**\n" + ", ".join(numeric_cols) + "\n\n"
            if categorical_cols:
                preview += "**分类型列：**\n" + ", ".join(categorical_cols) + "\n\n"
            
            self.chat_history = []
            return preview, gr.update(choices=columns), gr.update(choices=columns), gr.update(choices=columns)
            
        except Exception as e:
            return f"文件加载错误: {str(e)}", gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[])
    
    def chat_with_excel(self, message, history):
        """与Excel数据对话"""
        try:
            if self.current_df is None:
                return history + [{"role": "assistant", "content": "请先上传Excel文件"}]
            
            # 使用_build_data_context构建数据上下文
            data_context = self._build_data_context(message)
            
            # 调用大模型分析
            response = self.analyzer.chat.chat_completions(data_context)
            
            # 格式化响应
            formatted_response = self._format_response(response)
            
            # 更新对话历史
            self.chat_history.append({"role": "user", "content": message})
            self.chat_history.append({"role": "assistant", "content": formatted_response})
            
            return history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": formatted_response}
            ]
            
        except Exception as e:
            error_message = f"分析失败: {str(e)}"
            return history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": error_message}
            ]
    
    def _format_response(self, response):
        """格式化响应内容"""
        try:
            # 如果响应是JSON字符串，解析它
            if isinstance(response, str):
                try:
                    response_json = json.loads(response)
                except:
                    return response
            else:
                response_json = response
            
            # 提取回答内容
            if isinstance(response_json, dict):
                if 'Choices' in response_json and response_json['Choices']:
                    content = response_json['Choices'][0]['Message']['Content']
                    return content
                elif 'content' in response_json:
                    return response_json['content']
            
            # 如果无法解析，返回原始响应
            return str(response_json)
            
        except Exception as e:
            print(f"响应格式化错误: {e}")
            return str(response)
    
    def plot_chart(self, chart_type, x_col, y_col, z_col, title):
        """代理到 analyzer 的图表生成方法"""
        try:
            # 只传入必要的参数
            return self.analyzer.generate_chart(
                chart_type, 
                x_col, 
                y_col, 
                z_col,
                title
            )
        except Exception as e:
            return f"生成图表失败: {str(e)}", None
   
  
   
    def create_interface(self):
        """创建Gradio界面"""
        interface = gr.Blocks()
        
        with interface:
            gr.Markdown("""
            # Excel 智能分析助手
            
            上传Excel文件或图片，开始智能分析对话
            """)
            
            with gr.Tabs():
                # 数据分析标签页
                with gr.Tab("数据分析"):
                    with gr.Row():
                        # 左侧：文件上传和预览
                        with gr.Column(scale=1):
                            file_input = gr.File(
                                label="上传数据文件",
                                file_types=[".xlsx", ".xls", ".png", ".jpg", ".jpeg"]
                            )
                            data_preview = gr.Markdown(elem_classes="file-info")
                        
                        # 右侧：智能对话
                        with gr.Column(scale=1):
                            gr.Markdown("### 智能对话")
                            chatbot = gr.Chatbot(
                                elem_classes="chat-window",
                                height=400,
                                type="messages"
                            )
                            with gr.Row():
                                msg = gr.Textbox(
                                    placeholder="请输入你的问题...",
                                    lines=2
                                )
                                send = gr.Button("发送", variant="primary")
                
                    # 底部：数据分析报告
                    gr.Markdown("### 数据分析报告")
                    with gr.Row():
                        analyze_btn = gr.Button("生成分析报告", variant="primary")
                        export_btn = gr.Button("导出报告", variant="secondary")
                    analysis_report = gr.Markdown(elem_classes="file-info")
                
                # 数据可视化标签页
                with gr.Tab("数据可视化"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            chart_type = gr.Dropdown(
                                choices=list(self.available_charts.keys()),
                                label="选择图表类型",
                                value=list(self.available_charts.keys())[0]
                            )
                            x_column = gr.Dropdown(
                                label="X轴数据",
                                choices=[],
                                value=None,
                                allow_custom_value=False
                            )
                            y_column = gr.Dropdown(
                                label="Y轴数据",
                                choices=[],
                                value=None,
                                allow_custom_value=False
                            )
                            z_column = gr.Dropdown(
                                label="Z轴/大小数据",
                                choices=[],
                                value=None,
                                allow_custom_value=False,
                                visible=False
                            )
                            chart_title = gr.Textbox(
                                label="图表标题",
                                placeholder="输入图表标题..."
                            )
                            plot_btn = gr.Button("生成图表", variant="primary")
                        
                        with gr.Column(scale=1):
                            plot_output = gr.Image(label="图表预览")
                            plot_status = gr.Textbox(label="状态信息")
            
            # 添加图表类型变化事件
            def update_z_axis_visibility(chart_type):
                return gr.update(visible=chart_type in ["气泡图", "3D散点图"])
            
            chart_type.change(
                fn=update_z_axis_visibility,
                inputs=[chart_type],
                outputs=[z_column]
            )
            
            # 事件处理
            file_input.change(
                fn=self.load_excel,
                inputs=[file_input],
                outputs=[data_preview, x_column, y_column, z_column]  # 文件加载后更新下拉列表
            )
            
            analyze_btn.click(
                fn=self.analyzer.analyze_data,
                inputs=[],
                outputs=[analysis_report]
            )
            
            send.click(
                fn=self.chat_with_excel,
                inputs=[msg, chatbot],
                outputs=[chatbot]
            ).then(
                fn=lambda: "",
                outputs=[msg]
            )
            
            plot_btn.click(
                fn=self.plot_chart,
                inputs=[chart_type, x_column, y_column, z_column, chart_title],
                outputs=[plot_status, plot_output]
            )
        
        return interface
    
    def analyze_data(self):
        """生成数据分析报告"""
        if self.current_df is None:
            return "请先上传数据文件"
        
        report = "## 数据分析报告\n\n"
        
        # 基本统计
        report += "### 1. 基本统计信息\n"
        report += self.current_df.describe().to_markdown()
        report += "\n\n"
        
        # 缺失值分析
        missing = self.current_df.isnull().sum()
        if missing.any():
            report += "### 2. 缺失值分析\n"
            report += missing.to_markdown()
            report += "\n\n"
        
        # 相关性分析
        numeric_cols = self.current_df.select_dtypes(include=['number']).columns
        if len(numeric_cols) > 1:
            report += "### 3. 相关性分析\n"
            corr = self.current_df[numeric_cols].corr()
            report += corr.to_markdown()
            report += "\n\n"
        
        # 分布分析
        report += "### 4. 数据分布\n"
        for col in numeric_cols:
            stats = self.current_df[col].describe()
            report += f"#### {col}\n"
            report += f"- 均值：{stats['mean']:.2f}\n"
            report += f"- 中位数：{stats['50%']:.2f}\n"
            report += f"- 标准差：{stats['std']:.2f}\n"
            report += f"- 偏度：{self.current_df[col].skew():.2f}\n"
            report += f"- 峰度：{self.current_df[col].kurtosis():.2f}\n\n"
        
        return report
    
    def export_analysis(self):
        """导出分析结果"""
        try:
            if not self.chat_history:
                return "没有可导出的分析结果"
            
            # 生成导出文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"数据分析报告_{timestamp}.md"
            
            # 生成报告内容
            content = "# 数据分析报告\n\n"
            content += "## 基本信息\n"
            if self.current_df is not None:
                content += f"- 数据大小：{self.current_df.shape[0]}行 × {self.current_df.shape[1]}列\n"
                content += f"- 列名：{', '.join(self.current_df.columns)}\n\n"
            
            content += "## 分析历史\n\n"
            for msg in self.chat_history:
                if msg["role"] == "user":
                    content += f"### 问题：\n{msg['content']}\n\n"
                else:
                    content += f"### 回答：\n{msg['content']}\n\n"
            
            # 保存文件
            with open(filename, "w", encoding="utf-8") as f:
                f.write(content)
            
            return f"分析报告已导出到：{filename}"
            
        except Exception as e:
            return f"导出失败: {str(e)}"
    
    def _build_data_context(self, message):
        """构建数据上下文"""
        basic_stats = {
            'total_rows': self.current_df.shape[0],
            'total_cols': self.current_df.shape[1],
            'null_count': self.current_df.isnull().sum().sum(),
            'numeric_cols': self.current_df.select_dtypes(include=['int64', 'float64']).columns.tolist(),
            'categorical_cols': self.current_df.select_dtypes(include=['object']).columns.tolist()
        }
        
        # 使用模型判断问题类型
        question_analysis_prompt = f"""
        请分析以下问题的类型，判断是否需要深度分析。
        
        问题：{message}
        
        请只回答"简单"或"深入"。
        - 简单：如计算平均值、求和、查找最大最小值等基础操作。
        - 深入：如分析趋势、寻找关系、提供建议、评估情况等需要综合分析的问题
        """
        
        # 调用模型判断
        analysis_type = self.analyzer.chat.chat_completions(question_analysis_prompt)
        needs_deep_analysis = "深入" in self._format_response(analysis_type)
        
        # 基础数据信息
        data_context = f"""
        数据基本信息：
        - 数据规模：{basic_stats['total_rows']}行 × {basic_stats['total_cols']}列
        - 缺失值：{basic_stats['null_count']}
        - 数值列：{', '.join(basic_stats['numeric_cols'])}
        - 分类列：{', '.join(basic_stats['categorical_cols'])}
        
        数据内容：
        {self.current_df.to_string()}
        """
        
        # 根据模型判断结果添加分析要求
        if needs_deep_analysis:
            data_context += """
            分析要求：
            1. 数据洞察
            - 识别关键特征和模式
            - 发现数据中的趋势
            - 分析异常和特殊情况
            
            2. 统计分析
            - 提供相关统计指标
            - 解释数据分布特征
            - 分析相关性（如适用）
            
            3. 业务建议
            - 提供数据驱动的见解
            - 给出可行的改进建议
            - 指出潜在的机会
            
            4. 可视化建议
            - 推荐合适的图表类型
            - 建议关注的关键维度
            5. 改进建议
            - 指出数据质量问题
            - 建议额外需要的数据
            - 提供优化分析的方向
            
            """
            
        else:
            data_context += """
            回答要求：
            - 直接回答问题
            - 确保计算准确
            - 使用简洁的语言
            """
        
        data_context += f"\n用户问题：{message}"
        
        return data_context
    


if __name__ == "__main__":
    ui = ChatExcelUI()
    interface = ui.create_interface()
    interface.launch()